package com.ddd.core.wapper;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.text.AntPathMatcher;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HtmlUtil;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.util.List;
import java.util.Map;

/**
 * @Author: zhongwenbin
 * @Date: 2023/6/28 21:46
 * @Description: xss过滤器与request缓存包装器
 **/
public class XssRequestWrapper extends HttpServletRequestWrapper {

    private ByteArrayOutputStream data;
    private final List<String> excludes;
    private Map<String, String[]> parameterMap;

    public XssRequestWrapper(HttpServletRequest request, List<String> excludes) throws IOException {
        super(request);
        this.excludes = excludes;
    }

    /**
     * 过滤请求体中的html
     *
     * @return javax.servlet.ServletInputStream
     * @description
     * @author zhongwenbin
     * @date 2023/7/3 14:39
     */
    @Override
    public ServletInputStream getInputStream() throws IOException {
        if (ObjectUtil.isNotEmpty(data)) {
            return new CustomServletInputStream(data.toByteArray());
        }
        data = new ByteArrayOutputStream();
        IoUtil.copy(super.getInputStream(), data);
        // 是过滤的url路径
        if (handlerInclude()) {
            return new CustomServletInputStream(data.toByteArray());
        }
        String str = data.toString();
        if (StrUtil.isBlank(str)) {
            return new CustomServletInputStream(data.toByteArray());
        }
        // xss 过滤
        this.data.reset();
        data.write(HtmlUtil.cleanHtmlTag(str).getBytes());
        return new CustomServletInputStream(data.toByteArray());
    }

    @Override
    public String getParameter(String name) {
        String[] params = getParameterMap().get(name);
        if (ArrayUtil.isEmpty(params)) {
            return null;
        }
        return params[0];
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        if (CollUtil.isNotEmpty(this.parameterMap)) {
            return parameterMap;
        }
        this.parameterMap = super.getParameterMap();
        this.parameterMap.forEach((name, values) -> {
            if (ArrayUtil.isNotEmpty(values) && !handlerInclude()) {
                for (int i = 0; i < values.length; i++) {
                    values[i] = HtmlUtil.cleanHtmlTag(values[i]);
                }
            }
        });
        return parameterMap;
    }

    @Override
    public String[] getParameterValues(String name) {
        return getParameterMap().get(name);
    }

    /**
     * @return boolean
     * @description 是否是过滤的url路径
     * @author zhongwenbin
     * @date 2023/6/30 9:34
     */
    private boolean handlerInclude() {
        String url = ((HttpServletRequest) getRequest()).getServletPath();
        for (String exclude : this.excludes) {
            AntPathMatcher antPathMatcher = new AntPathMatcher();
            if (antPathMatcher.match(exclude, url)) {
                return true;
            }
        }
        return false;
    }

    private static class CustomServletInputStream extends ServletInputStream {
        public ByteArrayInputStream bis;


        public CustomServletInputStream(byte[] bytes) {
            bis = new ByteArrayInputStream(bytes);
        }

        @Override
        public boolean isFinished() {
            return false;
        }

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setReadListener(ReadListener listener) {

        }

        @Override
        public int read() throws IOException {
            return bis.read();
        }
    }
}